'''
与伪分布式任务执行器兼容的任务模块
数据集: CJEval_math
'''

import time
import os
import json
import logging
import sys, os, random
from pprint import pprint


# 导入根目录 (my_work 文件夹中)
sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))

# 导入大语言模型
from BM import Da_Mo_Xing_2
# 导入算法 !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
from simple_FE.simplified_FE import SimpleFE
# 导入: 执行模板, 执行模式匹配器, 评估模板（封闭问题）, 评估模式匹配器（封闭问题）
from youyi_evaluator import CarryMatcher_template, CarryMatcher, EvaluationMatcher_close_template, EvaluationMatcher_close

# 导入数据集
from youyi_data import read_CJEval_math_1


# 从环境变量获取节点ID
node_id = os.environ.get('NODE_ID', 'unknown')
# 获取节点日志器（单独运行的话就是创建）
logger = logging.getLogger(f"node_{node_id}")






def run_task(carry_model_name, optimize_model_name, evaluate_model_name):
    """
    执行任务
    
    参数:
        carry_model_name: 执行模型配置
        optimize_model_name: 优化模型配置
        evaluate_model_name: 评估模型配置
        
    返回:
        dict: 任务执行结果
    """
    
    # !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    algorithm_name = "FE"
    dataset_name = "CJEval_math"

    logger.info(f"初始化任务: algorithm={algorithm_name}, dataset={dataset_name}, carry={carry_model_name}, optimize={optimize_model_name}, evaluate={evaluate_model_name}")

    # 创建模型实例（算法 & 数据集 & 模型类型 & 模型名称）
    carry_model = Da_Mo_Xing_2(carry_model_name, 
                                  f"{algorithm_name} & {dataset_name} & carry_model & {carry_model_name}")
    optimize_model = Da_Mo_Xing_2(optimize_model_name,
                                  f"{algorithm_name} & {dataset_name} & optimize_model & {optimize_model_name}")
    evaluate_model = Da_Mo_Xing_2(evaluate_model_name,
                                  f"{algorithm_name} & {dataset_name} & evaluate_model & {evaluate_model_name}")

    # 创建实例 !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    FE = SimpleFE(optimize_model, optimize_model, carry_model)  # 分析模型、设计模型、测试模型
    
    
    # 读取 “CJEval中国初中考试任务”数据
    x_data, y_data = read_CJEval_math_1(500)
    # 将两个列表打包在一起
    combined = list(zip(x_data, y_data))
    # 打乱顺序
    #random.shuffle(combined)
    # 解包打乱后的列表
    shuffled_x_data, shuffled_y_data = zip(*combined)
    # 转换回列表形式
    x_data = list(shuffled_x_data)
    y_data = list(shuffled_y_data)
    # 划分训练集、验证集、测试集
    jie_xian_1 = int(len(x_data) * 0.6)
    jie_xian_2 = int(len(x_data) * 0.8)
    train_x, train_y = x_data[:jie_xian_1], y_data[:jie_xian_1]
    val_x, val_y = x_data[jie_xian_1: jie_xian_2], y_data[jie_xian_1: jie_xian_2]
    test_x, test_y = x_data[jie_xian_2:], y_data[jie_xian_2:]
    #print(train_x, train_y)
    #print(test_x, test_y)
    '''
    # 准备示例数据
    train_x = ["聪明", "直接", "正式地", "不受欢迎的", "减法"]
    train_y = ["愚蠢", "间接", "非正式地", "受欢迎的", "加法"]
    test_x = ["高", "热", "硬", "新"]
    test_y = ["低", "冷", "软", "旧"]
    '''
    

    # 执行模板（让 执行模型 根据 提示词、输入生成输出）包含[PROMPT], [INPUT], [OUTPUT]标记
    # [PROMPT]替换为当前评估的提示词，[INPUT]替换为当前评估对象的输入，[OUTPUT]替换为空，可以没有
    carry_template = CarryMatcher_template
    # 与 执行模版 配套的 执行模式匹配器
    carry_matcher = CarryMatcher()

    # LLM_mark
    # 评估模板（让 评估模型 根据 提示词、输入、预期输出、实际输出 生成 评估值、评估解释）
    # [PROMPT]替换为当前评估的提示词，[INPUT]替换为给执行模型的输入，[OUTPUT_R]替换为标准答案，[OUTPUT_P]替换为执行模型的实际输出
    evaluation_template = EvaluationMatcher_close_template
    # 评估模式匹配器
    evaluation_matcher = EvaluationMatcher_close()
    # 评估模型列表（有标准答案参考，一个评估专家） 
    evaluation_models = [evaluate_model]
    eval_config = {
            'evaluation_index': 'LLM_mark',              # 评估指标
            'carry_template': carry_template,            # 执行模版
            'carry_matcher': carry_matcher,              # 执行模式匹配器，与 执行模版 配套
            'carry_model': carry_model,                  # 执行模型

            'evaluation_template': evaluation_template,  # 评估模板
            'evaluation_matcher': evaluation_matcher,    # 评估模式匹配器，用于解析评估模型的返回
            'evaluation_models': evaluation_models,      # 评估模型（用于打分式评估，代理专家）, list, 元素为评估模型实例
            'task_kind': '封闭',                         # 问题类型, str, ['封闭', '开放']
            'consistency_judgment': None,                # str, ["cohen_kappa", "cohen_kappa_linear", "cohen_kappa_quadratic", "ICC", None], 一致性判断方法
            'consistency_threshold': None,               # float, 一致性判断方法的阈值，推荐[0.2, 0.2, 0.2, 0.4, None]
        }
    

    # 评估优化配置(不设置可填为: None)
    eval_optimized_config = {
        'samples_number': 5,  # 样本数量
        'if_back': False,     # 是否放回抽样（默认不放回）
        'c': 1.0,             # 探索参数
        'T': 'auto_2_2',      # 评估总轮数（auto_2_{λ}，基于动态阈值的早停法）
    }
    


    # !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    logger.info(f"执行任务（找到最佳提示词）")
    # 任务描述
    task_description = "计算数学问题"

    # 找到最佳提示词
    val_situation, best_prompt, val_avg_score = FE.find_best_prompt(
        (train_x, train_y),      # 训练集 (输入列表, 输出列表)
        (val_x, val_y),      # 验证集 (输入列表, 输出列表)
        eval_config,            # 评估配置
        eval_optimized_config,  # 评估优化配置
        None,                   # 初始提示词列表（None表示自动生成）
        5,                      # 示例样本数量
        task_description,       # 任务描述
        3,                      # 迭代次数
        5,                      # 种群数量
        2,                      # 亲代数量
        "wheel",                # 选择方法
        5,                      # 评估样本数量
        5,                      # 训练样本数量
        "random",               # 采样方法
        "base",                 # 反馈生成方法
        False                   # 是否使用知识库积累
    )

    


    logger.info(f"执行任务（测试集上评估）")
    # 在测试集上评估
    # 评估优化配置(不设置可填为: None)（测试的时候最好是全部测试）
    test_eval_optimized_config = {
        'samples_number': None, # 样本数量
        'if_back': False, # 是否放回抽样（默认不放回）
        'c': None, # 探索参数
        'T': None, # 评估总轮数
    }
    # !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    test_situation, test_results = FE.evaluate_prompt(
        [best_prompt], test_x, test_y,                                  # 基本参数
        eval_config, test_eval_optimized_config                       # 评估指标参数 与 batch-UCB 评估优化参数
    )
    #print('############# 测试结束，测试状态为：')
    #print(test_situation)
    #print('############# 测试结果：')
    #pprint(test_results)
    #print('############# 测试结果完毕 #############'
    
    # results [[prompt, combined_data, average_score]，，，，，]
    # combined_data: (x, y_real, y_pred_original, y_pred, 得分, 评估解释)
    test_avg_score = test_results[0][2]


    logger.info(f"任务完成")

    # 工作节点（本次任务）返回的结果
    # 验证状态、验证集上评估结果、最佳提示词、测试状态、测试集上评估结果、相关参数
    node_results = {
        "val_situation": val_situation,
        "val_avg_score": val_avg_score,
        "best_prompt": best_prompt,
        "test_situation": test_situation,
        "test_avg_score": test_avg_score,

        "algorithm_name": algorithm_name,
        "dataset_name": dataset_name,
        "carry_model_name": carry_model_name,
        "optimize_model_name": optimize_model_name,
        "evaluate_model_name": evaluate_model_name,
    }
    # 返回结果
    return node_results



if __name__ == "__main__":
    node_results = run_task("Qwen2.5-7B-Instruct", "Qwen2.5-7B-Instruct", "glm-4-flash")
    pprint(node_results)